"""
用于分析情绪，有两种分析模式
"""

from .text_flesh import words_set

__all__ = ['emotion_analyse', 'emotion_stat', 'freq']


def emotion_stat(words, emotion_words_list):
    """
    统计情绪词频

    :param words: 词语集
    :param emotion_words_list: 情绪词表
    :return: 情绪词频
    """

    emotion_num = [0, 0, 0, 0, 0]

    def emotion_count(words):
        for word in words:
            for i in range(5):
                if word in emotion_words_list[i]:
                    emotion_num[i] += 1

    emotion_count(words)

    return emotion_num


def freq(num_list, mode):
    """
    将情绪词频转换为分析结果。

    :param num_list: 词频列表，元素为各句子中五种情绪词的词频组成的列表
    :param mode: 分析模式，默认为比例模式，值为1则按词频最大优先分析
    :return: 情绪向量集
    """

    if mode == 1:
        m = max(num_list)
        num_list = [1 if x == m else 0 for x in num_list] if m != 0 else [
            0, 0, 0, 0, 0]
    s = sum(num_list)
    return [x / s for x in num_list] if s != 0 else num_list


def emotion_analyse(wordslist, mode=1):
    """
    分析文本情绪。
    输入文本分词结果，导入情绪词典并统计词频，最后进行分析。

    :param wordslist: 文本分词结果，列表中的每个元素为一句话的分词结果
    :param mode: 分析模式，默认为比例模式，值为1则按词频最大优先分析
    :return: 返回情绪向量集
    """

    anger = words_set("./data/anger.txt")
    disgust = words_set("./data/disgust.txt")
    fear = words_set("./data/fear.txt")
    joy = words_set("./data/joy.txt")
    sadness = words_set("./data/sadness.txt")
    emotion_words_list = [anger, disgust, fear, joy, sadness]

    emotion_list = []
    for words in wordslist:
        emotion_stat_result = emotion_stat(words, emotion_words_list)
        emotion_list.append(freq(emotion_stat_result, mode))

    return emotion_list
